{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "gpuClass": "standard",
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Embeddings index format for open data access\n",
        "\n",
        "[txtai](https://github.com/neuml/txtai) is an all-in-one embeddings database for semantic search, LLM orchestration and language model workflows.\n",
        "\n",
        "The main programming language with txtai is Python. A key tenet is that the underlying data in an embeddings index is accessible without txtai.\n",
        "\n",
        "This notebook will demonstrate this through a series of examples.\n"
      ],
      "metadata": {
        "id": "-xU9P9iSR-Cy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Install dependencies\n",
        "\n",
        "Install `txtai` and all dependencies."
      ],
      "metadata": {
        "id": "shlUi2kKS7KT"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "xEvX9vCpn4E0"
      },
      "outputs": [],
      "source": [
        "%%capture\n",
        "!pip install git+https://github.com/neuml/txtai#egg=txtai[graph] datasets sqlite-vec"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Load dataset\n",
        "\n",
        "This example will use the `chatgpt-prompts` dataset."
      ],
      "metadata": {
        "id": "408IyXzKFSiG"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from datasets import load_dataset\n",
        "\n",
        "dataset = load_dataset(\"fka/awesome-chatgpt-prompts\", split=\"train\")"
      ],
      "metadata": {
        "id": "IQ_ns6YvFRm1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Build an Embeddings index\n",
        "\n",
        "Let's first build an embeddings index using txtai."
      ],
      "metadata": {
        "id": "AtEdP7Utw3mk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from txtai import Embeddings\n",
        "\n",
        "embeddings = Embeddings()\n",
        "embeddings.index((x[\"act\"], x[\"prompt\"]) for x in dataset)\n",
        "embeddings.save(\"txtai-index\")"
      ],
      "metadata": {
        "id": "DPWrubv5oOn7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's take a look at the index that was created"
      ],
      "metadata": {
        "id": "0zrjOl6JtPI2"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!ls -l txtai-index\n",
        "!echo\n",
        "!file txtai-index/*"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "i7PN9itKtIyx",
        "outputId": "ac904afb-83a8-4c5f-cf62-f5340225b8ac"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total 268\n",
            "-rw-r--r-- 1 root root    342 Sep  6 15:21 config.json\n",
            "-rw-r--r-- 1 root root 262570 Sep  6 15:21 embeddings\n",
            "-rw-r--r-- 1 root root   2988 Sep  6 15:21 ids\n",
            "\n",
            "txtai-index/config.json: JSON data\n",
            "txtai-index/embeddings:  data\n",
            "txtai-index/ids:         data\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The txtai embeddings index format is documented [here](https://neuml.github.io/txtai/embeddings/format/). Looking at the files above, we have configuration, embeddings data and ids storage. Ids storage is only used when content is disabled.\n",
        "\n",
        "Let's inspect each file."
      ],
      "metadata": {
        "id": "hx8H0dpXtX5b"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "with open(\"txtai-index/config.json\") as f:\n",
        "  print(json.dumps(json.load(f), sort_keys=True, indent=2))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "h0yrQ8rqtnzh",
        "outputId": "3379ab18-1805-4ef1-b946-1a61d18522db"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{\n",
            "  \"backend\": \"faiss\",\n",
            "  \"build\": {\n",
            "    \"create\": \"2024-09-06T15:21:11Z\",\n",
            "    \"python\": \"3.10.12\",\n",
            "    \"settings\": {\n",
            "      \"components\": \"IDMap,Flat\"\n",
            "    },\n",
            "    \"system\": \"Linux (x86_64)\",\n",
            "    \"txtai\": \"7.5.0\"\n",
            "  },\n",
            "  \"dimensions\": 384,\n",
            "  \"offset\": 170,\n",
            "  \"path\": \"sentence-transformers/all-MiniLM-L6-v2\",\n",
            "  \"update\": \"2024-09-06T15:21:11Z\"\n",
            "}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import faiss\n",
        "\n",
        "index = faiss.read_index(\"txtai-index/embeddings\")\n",
        "print(f\"Total records {index.ntotal}\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-aqiqfqeuM5p",
        "outputId": "4b7856f0-ec8d-4960-a419-e88a14d988a8"
      },
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Total records 170\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import msgpack\n",
        "\n",
        "with open(\"txtai-index/ids\", \"rb\") as f:\n",
        "  print(msgpack.unpack(f)[5:10])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "I0ffDyJsunrW",
        "outputId": "d8c5072d-d181-46e1-fc1a-cc713a78e69f"
      },
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "['JavaScript Console', 'Excel Sheet', 'English Pronunciation Helper', 'Spoken English Teacher and Improver', 'Travel Guide']\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Each file can be read without txtai. [JSON](https://www.json.org/json-en.html), [MessagePack](https://msgpack.org/index.html) and [Faiss](https://github.com/facebookresearch/faiss) all have libraries in multiple programming languages."
      ],
      "metadata": {
        "id": "0e3dpAFuvP-_"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Embeddings index with SQLite\n",
        "\n",
        "In the next example, we'll use SQLite to store content and vectors courtesy of the [sqlite-vec](https://github.com/asg017/sqlite-vec) library."
      ],
      "metadata": {
        "id": "UUwk13mzwUTS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from txtai import Embeddings\n",
        "\n",
        "embeddings = Embeddings(content=True, backend=\"sqlite\")\n",
        "embeddings.index((x[\"act\"], x[\"prompt\"]) for x in dataset)\n",
        "embeddings.save(\"txtai-sqlite\")"
      ],
      "metadata": {
        "id": "Z80FWhuNwj14"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's once again explore the generated index files."
      ],
      "metadata": {
        "id": "Kxcm42rixAH0"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!ls -l txtai-sqlite\n",
        "!echo\n",
        "!file txtai-sqlite/*"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "PEHT6LqHw_lw",
        "outputId": "81d707dd-9760-40cd-e0b2-be43c26ba2d1"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total 1696\n",
            "-rw-r--r-- 1 root root     384 Sep  6 15:21 config.json\n",
            "-rw-r--r-- 1 root root  126976 Sep  6 15:21 documents\n",
            "-rw-r--r-- 1 root root 1605632 Sep  6 15:21 embeddings\n",
            "\n",
            "txtai-sqlite/config.json: JSON data\n",
            "txtai-sqlite/documents:   SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 31, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n",
            "txtai-sqlite/embeddings:  SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 392, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "This time note how there is a documents file with content stored in SQLite and a separate SQLite file for embeddings. Let's test it out."
      ],
      "metadata": {
        "id": "UXiKSG0JxLPo"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "embeddings.search(\"teacher\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VYMTXtUpxSLd",
        "outputId": "6259f3c2-6a76-434d-d40f-f11cb2e63c80"
      },
      "execution_count": 10,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[{'id': 'Math Teacher',\n",
              "  'text': 'I want you to act as a math teacher. I will provide some mathematical equations or concepts, and it will be your job to explain them in easy-to-understand terms. This could include providing step-by-step instructions for solving a problem, demonstrating various techniques with visuals or suggesting online resources for further study. My first request is \"I need help understanding how probability works.\"',\n",
              "  'score': 0.3421396017074585},\n",
              " {'id': 'Educational Content Creator',\n",
              "  'text': 'I want you to act as an educational content creator. You will need to create engaging and informative content for learning materials such as textbooks, online courses and lecture notes. My first suggestion request is \"I need help developing a lesson plan on renewable energy sources for high school students.\"',\n",
              "  'score': 0.3267676830291748},\n",
              " {'id': 'Philosophy Teacher',\n",
              "  'text': 'I want you to act as a philosophy teacher. I will provide some topics related to the study of philosophy, and it will be your job to explain these concepts in an easy-to-understand manner. This could include providing examples, posing questions or breaking down complex ideas into smaller pieces that are easier to comprehend. My first request is \"I need help understanding how different philosophical theories can be applied in everyday life.\"',\n",
              "  'score': 0.30780404806137085}]"
            ]
          },
          "metadata": {},
          "execution_count": 10
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The top N results as expected. Let's again inspect the files."
      ],
      "metadata": {
        "id": "R0zsQx37xjWz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "with open(\"txtai-sqlite/config.json\") as f:\n",
        "  print(json.dumps(json.load(f), sort_keys=True, indent=2))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "jDvH6dHrxqwf",
        "outputId": "dcce1647-0488-4f3e-dba0-667a932bad27"
      },
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "{\n",
            "  \"backend\": \"sqlite\",\n",
            "  \"build\": {\n",
            "    \"create\": \"2024-09-06T15:21:13Z\",\n",
            "    \"python\": \"3.10.12\",\n",
            "    \"settings\": {\n",
            "      \"sqlite\": \"3.37.2\",\n",
            "      \"sqlite-vec\": \"v0.1.1\"\n",
            "    },\n",
            "    \"system\": \"Linux (x86_64)\",\n",
            "    \"txtai\": \"7.5.0\"\n",
            "  },\n",
            "  \"content\": true,\n",
            "  \"dimensions\": 384,\n",
            "  \"offset\": 170,\n",
            "  \"path\": \"sentence-transformers/all-MiniLM-L6-v2\",\n",
            "  \"update\": \"2024-09-06T15:21:13Z\"\n",
            "}\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import sqlite3, sqlite_vec\n",
        "\n",
        "db = sqlite3.connect(\"txtai-sqlite/documents\")\n",
        "print(db.execute(\"SELECT COUNT(*) FROM sections\").fetchone()[0])\n",
        "\n",
        "db = sqlite3.connect(\"txtai-sqlite/embeddings\")\n",
        "db.enable_load_extension(True)\n",
        "sqlite_vec.load(db)\n",
        "print(db.execute(\"SELECT COUNT(*) FROM vectors\").fetchone()[0])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "AsalQ-uUxxO0",
        "outputId": "2a0e303e-e10e-424c-a952-e473d86cd2db"
      },
      "execution_count": 12,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "170\n",
            "170\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "As in the previous example, each file can be read without txtai. [JSON](https://www.json.org/json-en.html), [SQLite](https://www.sqlite.org/) and [sqlite-vec](https://github.com/asg017/sqlite-vec) all have libraries in multiple programming languages."
      ],
      "metadata": {
        "id": "Fo7XEUNny6s5"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Graph storage\n",
        "\n",
        "Starting with txtai 7.4, graphs are stored using MessagePack. The indexed file has a list of nodes and edges that can easily be imported."
      ],
      "metadata": {
        "id": "Ipu08WhL0j6p"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from txtai import Embeddings\n",
        "\n",
        "embeddings = Embeddings(content=True, backend=\"sqlite\", graph={\"approximate\": False})\n",
        "embeddings.index((x[\"act\"], x[\"prompt\"]) for x in dataset)\n",
        "embeddings.save(\"txtai-graph\")"
      ],
      "metadata": {
        "id": "cmFiae6j0wHz"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!ls -l txtai-graph\n",
        "!echo\n",
        "!file txtai-graph/*"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hnv82EAB059Q",
        "outputId": "454b263b-8733-4997-eb6e-7d585cfa32c6"
      },
      "execution_count": 14,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "total 1816\n",
            "-rw-r--r-- 1 root root     454 Sep  6 15:21 config.json\n",
            "-rw-r--r-- 1 root root  126976 Sep  6 15:21 documents\n",
            "-rw-r--r-- 1 root root 1605632 Sep  6 15:21 embeddings\n",
            "-rw-r--r-- 1 root root  119970 Sep  6 15:21 graph\n",
            "\n",
            "txtai-graph/config.json: JSON data\n",
            "txtai-graph/documents:   SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 31, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n",
            "txtai-graph/embeddings:  SQLite 3.x database, last written using SQLite version 3037002, file counter 1, database pages 392, cookie 0x1, schema 4, UTF-8, version-valid-for 1\n",
            "txtai-graph/graph:       data\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import msgpack\n",
        "\n",
        "with open(\"txtai-graph/graph\", \"rb\") as f:\n",
        "  data = msgpack.unpack(f)\n",
        "  print(data.keys())\n",
        "\n",
        "  for key in data:\n",
        "    if data[key]:\n",
        "      print(key, data[key][100])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Bx3Tzt3l08R3",
        "outputId": "662d157c-6912-4626-db64-d959c2719331"
      },
      "execution_count": 15,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "dict_keys(['nodes', 'edges', 'categories', 'topics'])\n",
            "nodes [100, {'id': 'Ascii Artist', 'text': 'I want you to act as an ascii artist. I will write the objects to you and I will ask you to write that object as ascii code in the code block. Write only ascii code. Do not explain about the object you wrote. I will say the objects in double quotes. My first object is \"cat\"'}]\n",
            "edges [5, 100, {'weight': 0.39010339975357056}]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Wrapping up\n",
        "\n",
        "This notebook gave an overview of the txtai embeddings index file format and how it supports open data access.\n",
        "\n",
        "While txtai can be used as an all-in-one embeddings database, it can also be used for only one part of the stack such as data ingestion. For example, it can be used to populate a Postgres or SQLite database for downstream use. The options are there."
      ],
      "metadata": {
        "id": "y7N-YZlR5S-0"
      }
    }
  ]
}